import csv
import pandas as pd
import numpy as np
from mayavi import mlab
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
from tvtk.api import tvtk
from scipy.interpolate import griddata,RBFInterpolator
from mpl_toolkits.mplot3d import Axes3D

def main():
    # 文件路径
    file_path = "/opt/star-xp-master/build/source/0-Simulation/extracted_data/event_81.csv"
    # 读取Deposits Info部分的数据
    deposits_data = []
    with open(file_path, 'r') as f:
        reader = csv.reader(f)
        in_deposits = False
        headers = None
        for row in reader:
            if row and row[0] == 'Deposits Info':
                in_deposits = True
                headers = next(reader)  # 下一行是列名
                continue
            if in_deposits and row:
                deposits_data.append(row)
    # 转换为DataFrame
    df = pd.DataFrame(deposits_data, columns=headers)
    # 提取需要的列并转换为数值
    pre_pos_x = df['pre_pos_x'].astype(float)
    pre_pos_y = df['pre_pos_y'].astype(float)
    pre_pos_z = df['pre_pos_z'].astype(float)
    post_pos_x = df['post_pos_x'].astype(float)
    post_pos_y = df['post_pos_y'].astype(float)
    post_pos_z = df['post_pos_z'].astype(float)
    edep = df['edep'].astype(float)*1e6

    # 计算中间点坐标
    mid_x = (pre_pos_x + post_pos_x) / 2*1000
    mid_y = (pre_pos_y + post_pos_y) / 2*1000
    mid_z = (pre_pos_z + post_pos_z) / 2*1000
    #
    list1 = []
    list2 = []
    list3 = []
    list4 = []
    for i in range(len(mid_x)):
        if edep[i]>1e-9 :
            list1.append(mid_x[i])
            list2.append(mid_y[i])
            list3.append(mid_z[i])
            list4.append(edep[i])
    mid_x = np.array(list1)
    mid_y = np.array(list2)
    mid_z = np.array(list3)
    edep = np.array(list4)
    #
    # Create a figure
    fig = plt.figure()
    ax = fig.add_subplot(111, projection='3d')

    # Plot 3D scatter
    ax.scatter(mid_x, mid_y, mid_z, c=edep)
    ax.set_title('3D Scatter Plot')

    # Show plot
    plt.show()
    # 归一化能量值到0-1
    min_edep = edep.min()
    max_edep = edep.max()
    normalized_edep = (edep - min_edep) / (max_edep - min_edep) if max_edep != min_edep else np.zeros_like(edep)
    #
    mid_xi=np.linspace(mid_x.min(), mid_x.max(), 50)
    mid_yi=np.linspace(mid_y.min(), mid_y.max(), 50)
    mid_zi=np.linspace(mid_z.min(), mid_z.max(), 50)
    # mid_xj, mid_yj, mid_zj = np.meshgrid(mid_xi, mid_yi, mid_zi)
    # edepj = griddata((mid_x, mid_y, mid_z), normalized_edep, (mid_xj, mid_yj, mid_zj), method='linear')
    # edepj[np.isnan(edepj)] = 0
    xgrid = np.mgrid[mid_x.min():mid_x.max():50j, mid_y.min():mid_y.max():50j, mid_z.min():mid_z.max():50j]
    xflat = xgrid.reshape(3, -1).T
    xobs = np.vstack([mid_x, mid_y, mid_z]).T
    yobs = normalized_edep
    yflat = RBFInterpolator(xobs, yobs, kernel='cubic')(xflat)
    ygrid = yflat.reshape(50, 50, 50)
    XI, YI, ZI = xgrid
    edepj = ygrid

    mlab.contour3d( XI, YI, ZI, edepj, contours=10, transparent=True)
    mlab.show()

if __name__ == "__main__":
    main()
    
